/**
 * Copyright 2019-2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "ops_store/sub_op_info_store.h"
#include <dirent.h>
#include "common/fe_log.h"
#include "common/aicore_util_types.h"
#include "common/string_utils.h"
#include "common/fe_type_utils.h"
#include "ops_store/ops_kernel_utils.h"
#include "ops_store/ops_kernel_error_codes.h"

namespace fe {
SubOpInfoStore::SubOpInfoStore(const FEOpsStoreInfo &ops_store_info)
    : init_flag_(false), ops_store_info_(ops_store_info) {}
SubOpInfoStore::~SubOpInfoStore() {}

Status SubOpInfoStore::Initialize(const std::string &engine_name) {
  if (init_flag_) {
    return SUCCESS;
  }
  init_flag_ = true;

  if (ops_store_info_.fe_ops_store_name.empty()) {
    FE_LOGW("The name of ops information library is empty.");
    return OP_STORE_CFG_NAME_EMPTY;
  }
  if (ops_store_info_.cfg_file_path.empty()) {
    FE_LOGW("The config file path of op information library[%s] is empty.", ops_store_info_.fe_ops_store_name.c_str());
    return OP_STORE_CFG_FILE_EMPTY;
  }

  std::string cfg_real_path = GetRealPath(ops_store_info_.cfg_file_path);
  if (cfg_real_path.empty()) {
    FE_LOGW("The config file[%s] of op information library[%s] is not existed. ",
            ops_store_info_.cfg_file_path.c_str(), ops_store_info_.fe_ops_store_name.c_str());
    return OP_STORE_CFG_FILE_NOT_EXIST;
  }

  FE_LOGD("Initialize %s SubOpsStore begin.", ops_store_info_.fe_ops_store_name.c_str());
  Status ret_value = LoadOpInfo(cfg_real_path);
  if (ret_value != SUCCESS) {
    REPORT_FE_ERROR("[GraphOpt][Init] Initialize sub op info library:[%s] LoadOpInfo failed! Config file_path:[%s]! \
                    Return value is [%u].", ops_store_info_.fe_ops_store_name.c_str(),
                    ops_store_info_.cfg_file_path.c_str(), ret_value);
    return OP_STORE_READ_CFG_FILE_FAILED;
  }

  ret_value = ConstructOpKernelInfo(engine_name);
  if (ret_value != SUCCESS) {
    return ret_value;
  }


  FE_LOGI("Initialize %s SubOpsStore finished.", ops_store_info_.fe_ops_store_name.c_str());
  return SUCCESS;
}

Status SubOpInfoStore::LoadOpInfo(const string &real_path) {
  vector<string> ops_cfg_files;
  DIR *dir;
  struct dirent *dirp = nullptr;
  char *file_suffix = const_cast<char *>(".json");
  dir = opendir(real_path.c_str());
  FE_CHECK(dir == nullptr, FE_LOGE("Fail to open directory %s.", real_path.c_str()), return FAILED);
  while ((dirp = readdir(dir)) != nullptr) {
    if (dirp->d_name[0] == '.') {
      continue;
    }
    if (strlen(dirp->d_name) <= strlen(file_suffix)) {
      continue;
    }
    if (0 == strcmp(&(dirp->d_name)[strlen(dirp->d_name) - strlen(file_suffix)], file_suffix)) {
      ops_cfg_files.push_back(real_path + "/" + dirp->d_name);
    }
  }
  closedir(dir);

  if (ops_cfg_files.empty()) {
    FE_LOGI("There is no json file in path %s.", real_path.c_str());
    return SUCCESS;
  }

  for (std::string json_file_path : ops_cfg_files) {
    if (LoadOpJsonFile(json_file_path) != SUCCESS) {
      REPORT_FE_ERROR("[GraphOpt][Init][LoadOpInfo] Fail to load json file[%s].", json_file_path.c_str());
      return FAILED;
    }
  }
  return SUCCESS;
}

Status SubOpInfoStore::LoadOpJsonFile(const std::string &json_file_path) {
  vector<string> op_type_vec;
  nlohmann::json op_json_file;
  FE_LOGD("Begin to load json file[%s].", json_file_path.c_str());
  if (ReadJsonObject(json_file_path, op_json_file) != SUCCESS) {
    REPORT_FE_ERROR("[GraphOpt][Init][LoadOpJsonFile] ReadJsonObject in %s failed.", json_file_path.c_str());
    return FAILED;
  }
  try {
    if (!op_json_file.is_object()) {
      REPORT_FE_ERROR("[GraphOpt][Init][LoadOpJsonFile] Top level of json file should be object, actually is %s.",
                      GetJsonObjectType(op_json_file).c_str());
      return OP_SUB_STORE_ILLEGAL_JSON;
    }
    for (auto &elem : op_json_file.items()) {
      string op_type_temp = elem.key();
      op_type_vec.push_back(StringUtils::Trim(op_type_temp));
    }
    for (auto &op_type : op_type_vec) {
      OpContent op_content;
      op_content.op_type_ = op_type;
      if (!op_json_file[op_type].is_object()) {
        REPORT_FE_ERROR("[GraphOpt][Init][LoadOpJsonFile] Second level of json file should be object, actually is %s.",
                GetJsonObjectType(op_json_file[op_type]).c_str());
        return OP_SUB_STORE_ILLEGAL_JSON;
      }
      for (auto &elem_out : op_json_file[op_type].items()) {
        map<string, string> map_temp;
        if (!op_json_file[op_type][elem_out.key()].is_object()) {
          REPORT_FE_ERROR("[GraphOpt][Init][LoadOpJsonFile] Third level of json file should be object, actually is %s.",
                  GetJsonObjectType(op_json_file[op_type][elem_out.key()]).c_str());
          return OP_SUB_STORE_ILLEGAL_JSON;
        }
        for (auto &elem_inner : op_json_file[op_type][elem_out.key()].items()) {
          string key_inner = elem_inner.key();
          string value_inner = elem_inner.value();
          map_temp.emplace(std::make_pair(StringUtils::Trim(key_inner), StringUtils::Trim(value_inner)));
        }
        string key_out = elem_out.key();
        op_content.map_kernel_info_.emplace(std::make_pair(StringUtils::Trim(key_out), map_temp));
      }
      op_content_map_.emplace(std::make_pair(op_type, op_content));
    }
  } catch (const std::exception &e) {
    REPORT_FE_ERROR("[GraphOpt][Init][LoadOpJsonFile] Fail to convfile[%s] to Json.Error message is %s",
                    json_file_path.c_str(), e.what());
    return OP_SUB_STORE_ILLEGAL_JSON;
  }
  return SUCCESS;
}

Status SubOpInfoStore::ConstructOpKernelInfo(const std::string &engine_name) {
  OpKernelInfoConstructor op_kernel_info_constructor;
  for (auto &op : op_content_map_) {
    OpKernelInfoPtr op_kernel_info_ptr = nullptr;
    FE_MAKE_SHARED(op_kernel_info_ptr = std::make_shared<OpKernelInfo>(op.first), return FAILED);
    FE_CHECK_NOTNULL(op_kernel_info_ptr);

    op_kernel_info_ptr->SetImplType(ops_store_info_.op_impl_type);
    /* Here shared ptr and map/vector will destruct automatically, so not necessary
       to run final for those op_kernel_info which fails to initialize. */
    if (op_kernel_info_constructor.InitializeOpKernelInfo(engine_name, op.second, op_kernel_info_ptr) != SUCCESS) {
      REPORT_FE_ERROR("[GraphOpt][Init][ConstructOpKernelInfo] opKernelInfo %s initialize failed.", op.first.c_str());
      return FAILED;
    }

    op_kernel_info_map_.emplace(std::make_pair(op.first, op_kernel_info_ptr));
  }
  return SUCCESS;
}

Status SubOpInfoStore::Finalize() {
  if (!init_flag_) {
    return SUCCESS;
  }
  OpKernelInfoConstructor op_kernel_info_constructor;
  for (auto &el : this->op_kernel_info_map_) {
    if (el.second == nullptr) {
      FE_LOGW("OpKernelInfo pointer[%s] in op_kernel_info_map_ should not be nullptr.",
              el.first.c_str());
      continue;
    }
    op_kernel_info_constructor.FinalizeOpKernelInfo(el.second);
  }
  this->op_kernel_info_map_.clear();
  init_flag_ = false;
  return SUCCESS;
}

const std::map<std::string, OpKernelInfoPtr>& SubOpInfoStore::GetAllOpKernels() {
  return this->op_kernel_info_map_;
}

OpKernelInfoPtr SubOpInfoStore::GetOpKernelByOpType(const std::string &op_type) {
  auto iter = op_kernel_info_map_.find(op_type);
  if (iter == op_kernel_info_map_.end()) {
    FE_LOGD("Op Type[%s] is not existed in op information library[%s].",
            op_type.c_str(), this->ops_store_info_.fe_ops_store_name.c_str());
    return nullptr;
  }
  return iter->second;
}

Status SubOpInfoStore::GetOpContentByOpType(const std::string &op_type, OpContent &op_content) {
  auto iter = op_content_map_.find(op_type);
  if (iter == op_content_map_.end()) {
    FE_LOGD("Op Type[%s] is not found in op information library[%s].",
            op_type.c_str(), this->ops_store_info_.fe_ops_store_name.c_str());
    return FAILED;
  }
  op_content = iter->second;
  return SUCCESS;
}

Status SubOpInfoStore::SetOpContent(const OpContent &op_content) {
  if (op_content.op_type_.empty()) {
    REPORT_FE_ERROR("[GraphOpt][SetDynamicinfo][SetOpContent] Op type of op_content is empty.");
    return FAILED;
  }
  op_content_map_[op_content.op_type_] = op_content;
  return SUCCESS;
}

const std::string& SubOpInfoStore::GetOpsStoreName() const {
  return this->ops_store_info_.fe_ops_store_name;
}

const OpImplType& SubOpInfoStore::GetOpImplType() const {
  return this->ops_store_info_.op_impl_type;
}

}  // namespace fe
